import numpy as np
from collections import Counter

class KNN:
    def __init__(self, k=3):
        self.k = k
        self.X_train = None
        self.y_train = None
    
    def fit(self, X, y):
        self.X_train = X
        self.y_train = y
    
    def predict(self, X):
        predictions = []
        for x in X:
            # 计算当前样本与所有训练样本的欧式距离
            distances = []
            for i, x_train in enumerate(self.X_train):
                dist = np.sqrt(np.sum((x - x_train) ** 2))
                distances.append((dist, i))
            
            # 按距离排序，取前k个
            distances.sort(key=lambda x: x[0])
            k_nearest = [self.y_train[i] for _, i in distances[:self.k]]
            
            # 投票选出最多的类别
            most_common = Counter(k_nearest).most_common(1)
            predictions.append(most_common[0][0])
        
        return np.array(predictions)

# 测试代码
if __name__ == "__main__":
    # 创建示例数据
    X_train = np.array([[1, 2], [2, 3], [3, 1], [6, 5], [7, 7], [8, 6]])
    y_train = np.array([0, 0, 0, 1, 1, 1])
    
    X_test = np.array([[2, 2], [7, 6]])
    
    # 使用KNN
    knn = KNN(k=3)
    knn.fit(X_train, y_train)
    predictions = knn.predict(X_test)
    
    print("训练数据:", X_train)
    print("训练标签:", y_train)
    print("测试数据:", X_test)
    print("预测结果:", predictions)